
from llama_index.core.data_structs import Node
from llama_index.core import QueryBundle
from llama_index.core.postprocessor.types import BaseNodePostprocessor
from llama_index.core.schema import NodeWithScore
from typing import List,Optional

from treesearch import *


class TreeSearchReranker(BaseNodePostprocessor):
    
    extor: RelaSubjectExtor
    depth: Optional[int]

    def __init__(self,related_dirs:list[str],depth:int=None):
        super().__init__(
            extor=RelaSubjectExtor(related_dirs), depth=depth
        )

    def _postprocess_nodes(
        self, nodes: List[NodeWithScore], query_bundle: Optional[QueryBundle]
    ) -> List[NodeWithScore]:
        result: List[NodeWithScore] = []
        visSet = set()
        # 遍历所有的评分节点
        for node in nodes:
            file_path = node.node.metadata.get("file_path",None)
            result.append(node)
            visSet.add(file_path)
            # TODO: 目前只对dirctor处理
            if "director" not in file_path:
                continue
            if file_path != None and os.path.exists(file_path):
                outDict = self.extor.SearchInTree(file_path,self.depth)
            # 组装Node
            for title,docNode in outDict.items():
                # 检查重复性
                if docNode.AttrFilePath() in visSet:
                    continue
                # 构建节点
                newNode = Node(text=docNode.LoadContent())
                newNode.metadata["document_title"] = title
                # 组装NodeWithScore
                result.append(NodeWithScore(node=newNode,score=node.score))
                # vis 标记
                visSet.add(docNode.AttrFilePath())
        print("selected nodes with reranker: \n",result,"\n\n")
        return result

